from typing import Literal
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from src.agentic.config.VectorStore import VectorStore
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class RoutingLogic:

    def __init__(self, llm_model, embed_model, collection_prefix, document_sources):
        self.llm_model = llm_model
        self.embed_model = embed_model
        self.collection_prefix = collection_prefix
        self.document_sources = document_sources

    def invoke(self, query):
        logger.info(f"RoutingLogic invoke query: {query}")

        datasource_list = []
        datasource_descs = []
        for source in self.document_sources:
            datasource_list.append(self.collection_prefix + source.library_number)
            datasource_descs.append({ source.library_number: source.document_summary })
        logger.info(f"RoutingLogic invoke datasource_list: {datasource_list}")

        class DataSource(BaseModel):
            datasource: Literal[*datasource_list] = Field(description = "根据用户问题，选择最相关的数据源来回答问题")

        datasource_prompt = f"""
            您擅长将用户问题转接到合适的数据源。
            根据问题涉及的编程语言，将其路由到相关数据源。
            数据源列表： {datasource_descs}
            用户问题：{query}
        """
        datasource_result = self.llm_model.with_structured_output(DataSource).invoke(datasource_prompt)
        logger.info(f"RoutingLogic invoke datasource_result: {datasource_result}")

        vector_store = VectorStore().new_vector_store(self.embed_model, datasource_result)
        retriever = vector_store.as_retriever(search_kwargs = { "k": 3 })
        retrieve_docs = retriever.invoke(query)

        template = """
            请基于以下上下文内容回答问题：
            {context}
            
            问题：{question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        chain_result = chain.invoke({ "context": retrieve_docs, "question": query })
        logger.info(f"RoutingLogic invoke chain_result len: {len(chain_result)}")
        return { "retrieve_docs": retrieve_docs, "chain_result": chain_result }
